import os
import json
import argparse
import torch
import dataloaders
import models
import inspect
import math
from utils import losses
from utils import Logger
from utils.torchsummary import summary
from trainer import Trainer
import torch.onnx
import numpy as np
from torch.autograd import Variable
from torchvision.models import AlexNet
from torchviz import make_dot

config = 'config.json'


def get_instance(module, name, config, *args):
    # GET THE CORRESPONDING CLASS / FCT
    # 从module(dataloaders)中找到config[name]['type'] (class VOC();class PSPNet()),将config[name]['args']的键值对作为参数的键值对传入,args为按顺序传入
    return getattr(module, config[name]['type'])(*args, **config[name]['args'])


def main(config, resume):
    train_logger = Logger()

    # 从config.json中获取参数，创建dataloaders中相应的类并传入相应参数创建获取数据的模型，从models中获取分类模型，从utils.losses中获取loss模型
    # DATA LOADERS
    train_loader = get_instance(dataloaders, 'train_loader', config)
    val_loader = get_instance(dataloaders, 'val_loader', config)

    # MODEL
    model = get_instance(models, 'arch', config, train_loader.dataset.num_classes)
    print(f'\n{model}\n')

    # # 模型可视化
    # x = torch.rand(8, 3, 512, 512)
    # y = model(x)
    # g = make_dot(y)
    # g.render('model_output', view=False)

    # LOSS
    # 根据utils.losses中定义的类，创建config['loss']类，并传入相应参数
    loss = getattr(losses, config['loss'])(ignore_index=config['ignore_index'])

    # TRAINING
    trainer = Trainer(
        model=model,
        loss=loss,
        resume=resume,
        config=config,
        train_loader=train_loader,
        val_loader=val_loader,
        train_logger=train_logger)

    trainer.train()


if __name__ == '__main__':
    # PARSE THE ARGS
    parser = argparse.ArgumentParser(description='PyTorch Training')
    parser.add_argument('-c', '--config', default='config.json', type=str,
                        help='Path to the config file (default: config.json)')
    parser.add_argument('-r', '--resume', default=None, type=str,
                        help='Path to the .pth model checkpoint to resume training')
    parser.add_argument('-d', '--device', default=None, type=str,
                        help='indices of GPUs to enable (default: all)')
    args = parser.parse_args()  # args.config == 'config.json'  args.device == None    args.resume == None

    config = json.load(open(args.config))
    if args.resume:
        config = torch.load(args.resume)['config']
    if args.device:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device
    
    main(config, args.resume)
